package com.google.code.securityflow.graphviz

// TODO: Low Priority: Reimplement graphviz
object Graphviz {

  /*
   * AST-related
   */
  def toDot(tree: GraphvizTree): String = {
    "digraph {\n" + vertices(tree) + "n" + edges(tree) + "}\n";
  }

  def vertices(tree: GraphvizTree): String = {
    tree match {
      case GraphvizTreeLeaf(id, label) => {
        "\tnode" + id + " [label=\"" + label + "\",shape=ellipse];\n"
      }
      case GraphvizTreeNode(id, label, children) => {
        "\tnode" + id + " [label=\"" + label + "\",shape=rectangle];\n" + children.foldLeft("")((s, t) => s + vertices(t));
      }
    }
  }

  def edges(tree: GraphvizTree): String = {
    val edges = tree.toList.map(t => {
      t match {
        case GraphvizTreeLeaf(id, label) => "";
        case GraphvizTreeNode(id, label, children) => {
          children.foldLeft("")((acc, c) => acc + "\tnode" + id + " -> " + "node" + c.id + ";\n");
        }
      }
    })
    edges.foldLeft("")((s, acc) => s + acc);
  }

  /*
  * CFG-related
  */
  def toDot(clusters: Set[GraphvizCluster]): String = {
    "digraph {\n\trankdir=\"TD\";\n" + clusters.foldLeft("")((s, c) => s + build(c)) + "}\n";
  }

  def build(cluster: GraphvizCluster): String = {
    "\tsubgraph cluster" + cluster.id + " {\n" + "\t\tlabel=\"" + cluster.label + "\"\n" +
      vertices(cluster.nodes) + "n" +
      edges(cluster.nodes) + "\t}\n";
  }

  def vertices(nodes: Set[GraphvizMultiNode]) = {
    val sb = new StringBuilder();
    for (node <- nodes) {
      sb.append("\t\tblock" + node.id + " [label=\"" + "{" + node.labels.mkString("|") + "}" + "\",shape=record];\n");
    }
    sb.toString();
  }

  def edges(nodes: Set[GraphvizMultiNode]) = {
    val sb = new StringBuilder();
    for (source <- nodes; target <- source.successors) {
      sb.append("\t\tblock" + source.id + " -> " + "block" + target.id + " [headport=n,tailport=s,headlabel=\"" + "" * 6 + target.id + "\"];\n")
    }
    sb.toString();
  }

}

/*
* AST-related
*/
abstract sealed class GraphvizTree(val id: Int, val label: String) {
  def toList: List[GraphvizTree];
}

case class GraphvizTreeLeaf(override val id: Int, override val label: String) extends GraphvizTree(id, label) {
  override def toList: List[GraphvizTree] = List(this);
}

case class GraphvizTreeNode(override val id: Int, override val label: String, children: List[GraphvizTree]) extends GraphvizTree(id, label) {
  override def toList: List[GraphvizTree] = List(this) ++ children.flatMap(c => c.toList);
}

/*
* CFG-related
*/
case class GraphvizCluster(id: Int, label: String, nodes: Set[GraphvizMultiNode]) {

}

case class GraphvizMultiNode(id: Int, labels: List[String], var successors: Set[GraphvizMultiNode]) {

  def addSuccessor(successor: GraphvizMultiNode) {
    successors = successors + successor;
  }

  override def equals(o: Any) = {
    o match {
      case that: GraphvizMultiNode => this eq that;
      case _ => false;
    }
  };

  override def hashCode() = id.hashCode + labels.hashCode;

}




